import os
import re
import gradio as gr
from langchain.chains.combine_documents import create_stuff_documents_chain
from langchain.chains.retrieval import create_retrieval_chain
from langchain_community.embeddings import ModelScopeEmbeddings
from langchain_community.vectorstores import FAISS
from langchain_core.prompts import ChatPromptTemplate
from langchain_openai import ChatOpenAI
from langchain_text_splitters import CharacterTextSplitter
from pydantic.v1 import SecretStr
from pathlib import Path

# 使用绝对路径（自动适配操作系统）
DB_DIR = Path(__file__).parent / "faiss_db"  # 自动处理路径分隔符
print(f"数据库路径：{DB_DIR}")
EMBEDDINGS = ModelScopeEmbeddings(model_id="iic/nlp_gte_sentence-embedding_chinese-large")

# 全局链对象
CHAIN = None

# 一、构建向量数据库
def save_vectors_db():
    if DB_DIR.exists():
        print(f"数据库已存在于: {DB_DIR}")
    else:
        print('没有找到数据库')
        os.mkdir(DB_DIR)
    with open("sales_rag_data.txt", encoding="utf-8") as f:
        content = f.read()
    textsplitter = CharacterTextSplitter(
        separator=r"(?m)^Interaction_\d+\n",
        is_separator_regex=True,
        chunk_size=250,
        chunk_overlap=0,
        length_function=len
    )
    docs = textsplitter.create_documents([content])
    print(f'分块长度：{len(docs)}')
    db = FAISS.from_documents(docs, EMBEDDINGS)
    db.save_local(DB_DIR)

# 二、初始化检索链
def init_chain():
    db = FAISS.load_local(DB_DIR, EMBEDDINGS, allow_dangerous_deserialization=True)
    system_prompt = '''
    你是一个智能家居销售问答助手，根据下面检索到的上下文片段来回答问题，如果你不知道答案，就说“你不知道”。
    最多使用三句话进行回答，保持回答简洁。\n
    {context}
    '''
    prompt_template = ChatPromptTemplate(
        [
            ('system', system_prompt),
            ('human', '{input}')
        ]
    )
    retriever = db.as_retriever(search_type="similarity_score_threshold", search_kwargs={"score_threshold": 0.5})
    model = ChatOpenAI(
        model='deepseek-chat',
        temperature=0.2,
        api_key="***",
        base_url="https://api.deepseek.com"
    )
    chain1 = create_stuff_documents_chain(llm=model, prompt=prompt_template)
    chain = create_retrieval_chain(retriever, chain1)
    return chain

# 三、Gradio 接口处理函数
def query_sales_bot(user_input):
    global CHAIN
    if not user_input.strip():
        return "请输入有效的问题。"
    try:
        result = CHAIN.invoke({"input": user_input})
        return result['answer']
    except Exception as e:
        return f"发生错误：{str(e)}"

# 四、Gradio Web 界面
def create_gradio_interface():
    with gr.Blocks(title="智能家居销售问答机器人") as demo:
        gr.Markdown("# 智能家居销售问答机器人")
        gr.Markdown("输入您的问题，机器人将根据产品数据回答（最多三句话，简洁明了）。")
        
        with gr.Row():
            with gr.Column(scale=3):
                input_text = gr.Textbox(label="您的问题", placeholder="例如：智能灯泡的寿命有多长？")
            with gr.Column(scale=1):
                submit_button = gr.Button("提交")
        
        output_text = gr.Textbox(label="机器人回答", interactive=False)
        
        submit_button.click(
            fn=query_sales_bot,
            inputs=input_text,
            outputs=output_text
        )
    
    return demo

if __name__ == '__main__':
    # 初始化数据库和链
    save_vectors_db()
    CHAIN = init_chain()  # 直接赋值，无需 global 声明
    
    # 启动 Gradio 界面
    demo = create_gradio_interface()
    demo.launch()